/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#ifndef LINAL_HIER_HXX
#define LINAL_HIER_HXX

#define LINAL_ELEM_FLA(name, row, col) linal::elem<FLA_Obj>(buf ## name , cs ## name , 1 , row, col)

namespace linal {
  typedef struct FLA_Obj_view FLA_Obj;
  typedef class  Matrix_*     Matrix;
  typedef class  Flat_*       Flat;
  typedef class  Hier_*       Hier;
  /*!
    Hier_ class only use one-level blocking.
  */
  class Hier_ : public Matrix_ {
  private:
  protected:
  public:
    Hier_() : Matrix_() { }
    Hier_(FLA_Obj obj) : Matrix_(obj) { }
    Hier_(const Hier_&b) : Matrix_(b) { }

    virtual ~Hier_() { this->free(); }

    // ---------------------------------------------------------
    // ** Create / Free
    // Create Hier_ matrix of matrices with size (m,n)
    void create(Int_ type, Int_ m, Int_ n) {
      Matrix_::create(type, FLA_MATRIX, m, n);
    }

    // Create Hier_ matrix with block size (mb, nb) for the matrix size (m,n)
    void create(Int_ type, 
                Int_ offm, Int_ offn, Int_ m, Int_ n, 
                Int_ mb, Int_ nb) {
      if (!m || !n) return;

      LINAL_ERROR( !this->is_created(),         LINAL_ERROR_OBJ_CREATED );
      LINAL_ERROR( check_all_scalar_type(type), LINAL_ERROR_CHECK_SCALAR ); 
      
      std::vector< Int_ > offm_list, offn_list;
      Int_ offs, mm,nn;
      FLA_Obj *buffer;

      // calculate the size of Hier_ matrix
      mm = m/mb + (m%mb > 0);
      nn = n/nb + (n%nb > 0);

      // create                                                     
      FLA_Obj_create_ext(type, FLA_MATRIX, mm, nn, 0, 0, 0, 0,
			 &this->_fla);

      buffer = (FLA_Obj*)FLA_Obj_buffer_at_view(this->_fla);

      // calculate offset values
      offm_list.reserve(mm+1);
      offs = offm;
      for (Int_ i=0;i<mm;++i) { offm_list.push_back(offs); offs += mb; }
      offm_list.push_back(offm + m);

      offn_list.reserve(nn+1);
      offs = offn;
      for (Int_ j=0;j<nn;++j) { offn_list.push_back(offs); offs += nb; }
      offn_list.push_back(offn + n);

      // assign block matrices to the matrix objects in the Hier_ buffer
      for (Int_ j=0;j<nn;j++) {
	for (Int_ i=0;i<mm;i++) {
	  FLA_Obj *part = &buffer[ j*mm + i ];

	  part->offm = offm_list.at(i);
	  part->offn = offn_list.at(j);

	  part->m    = offm_list.at(i+1) - offm_list.at(i);
	  part->n    = offn_list.at(j+1) - offn_list.at(j);

	  part->base = NULL;
	}
      }
      
      // this is created
      this->set_created(1);
    }

    // Create Hier_ matrix with block size (mb, nb) for the matrix size (m,n)
    void create(Int_ type, Int_ m, Int_ n, Int_ mb, Int_ nb) {
      this->create(type, 0, 0, m, n, mb, nb);
    }

    // Create Hier_ matrix with block size (mb, nb) for the Flat_ matrix.
    void create(Flat_ F, Int_ mb, Int_ nb) {
      if (!F.get_m() || !F.get_n()) return;
      this->create(F.get_datatype(), 
                   F.get_offm(), F.get_offn(), 
                   F.get_m(),    F.get_n(), 
                   mb, nb);

      void *buffer = this->get_buffer();
      Int_ cs = this->get_cs(), rs = this->get_rs();

      // base share
      for (Int_ j=0;j<this->get_n();++j) 
	for (Int_ i=0;i<this->get_m();++i) {
          elem<FLA_Obj>(buffer, cs, rs, i, j).base = (~F).base;
        }
    }

    // Create Hier_ matrix with merging four 2d partitioned Hier_ submatrices 
    void create_2x2(Hier_ ATL, Hier_ ATR,
                    Hier_ ABL, Hier_ ABR) {

      this->create(ATL.get_datatype(), 
                   ATL.get_m()+ABL.get_m(),
                   ATL.get_n()+ATR.get_n());

      void *buffer_ME  = this->get_buffer();
      void *buffer_TL  = ATL.get_buffer();
      void *buffer_BL  = ABL.get_buffer();
      void *buffer_TR  = ATR.get_buffer();
      void *buffer_BR  = ABR.get_buffer();

      Int_ cs_ME = this->get_cs();
      Int_ cs_TL  = ATL.get_cs();
      Int_ cs_BL  = ABL.get_cs();
      Int_ cs_TR  = ATR.get_cs();
      Int_ cs_BR  = ABR.get_cs();

      Int_ offm = ATL.get_m(), offn = ATL.get_n();
      
      // ATL
      for (Int_ j=0;j<ATL.get_n();j++) 
	for (Int_ i=0;i<ATL.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, i, j) = 
            elem<FLA_Obj>(buffer_TL, cs_TL, 1, i, j);
            
      // ATR
      for (Int_ j=0;j<ATR.get_n();j++) 
	for (Int_ i=0;i<ATR.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, i, offn+j) = 
            elem<FLA_Obj>(buffer_TR, cs_TR, 1, i, j);
      
      // ABL
      for (Int_ j=0;j<ABL.get_n();j++) 
	for (Int_ i=0;i<ABL.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, offm+i, j) = 
            elem<FLA_Obj>(buffer_BL, cs_BL, 1, i, j);
      
      // ABR
      for (Int_ j=0;j<ABR.get_n();j++) 
	for (Int_ i=0;i<ABR.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, offm+i, offn+j) =
            elem<FLA_Obj>(buffer_BR, cs_BR, 1, i, j);
    }
    
    void create_2x1(Hier_ ATL, 
                    Hier_ ABL) {

      this->create(ATL.get_datatype(), 
                   ATL.get_m()+ABL.get_m(), 
                   ATL.get_n());
      
      void *buffer_ME = this->get_buffer();
      void *buffer_TL = ATL.get_buffer();
      void *buffer_BL = ABL.get_buffer();

      Int_ cs_ME = this->get_cs();
      Int_ cs_TL = ATL.get_cs();
      Int_ cs_BL = ABL.get_cs();

      Int_ offm = ATL.get_m();

      // ATL
      for (Int_ j=0;j<ATL.get_n();j++) 
	for (Int_ i=0;i<ATL.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, i, j) = 
            elem<FLA_Obj>(buffer_TL, cs_TL, 1, i, j);

      // ABL
      for (Int_ j=0;j<ABL.get_n();j++) 
	for (Int_ i=0;i<ABL.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, offm+i, j) = 
            elem<FLA_Obj>(buffer_BL, cs_BL, 1, i, j);
    }

    void create_1x2(Hier_ ATL, Hier_ ATR) {
                              
      this->create(ATL.get_datatype(), 
                   ATL.get_m(),
                   ATL.get_n()+ATR.get_n());
      
      void *buffer_ME = this->get_buffer();
      void *buffer_TL = ATL.get_buffer();
      void *buffer_TR = ATR.get_buffer();

      Int_ cs_ME = this->get_cs();
      Int_ cs_TL = ATL.get_cs();
      Int_ cs_TR = ATR.get_cs();

      Int_ offn = ATL.get_n();

      // ATL
      for (Int_ j=0;j<ATL.get_n();j++) 
	for (Int_ i=0;i<ATL.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, i, j) = 
            elem<FLA_Obj>(buffer_TL, cs_TL, 1, i, j);
      
      // ATR
      for (Int_ j=0;j<ATR.get_n();j++) 
	for (Int_ i=0;i<ATR.get_m();i++) 
          elem<FLA_Obj>(buffer_ME, cs_ME, 1, i, offn+j) = 
            elem<FLA_Obj>(buffer_TR, cs_TR, 1, i, j);
    }

    void wrap(FLA_Obj &obj)  {
      LINAL_ERROR( obj.base != NULL && obj.base->elemtype == FLA_MATRIX,
                   LINAL_ERROR_CHECK_MATRIX_ELEMTYPE );
      Matrix_::wrap(obj);
    }

    // Coarse the block matrix by merging the block with given band bm and bn.
    void coarse(Hier_ &A, Int_ bm, Int_ bn) {
      // LINAL_ERROR ... 
      A.set_view( (this->get_m()/bm) + (this->get_m()%bm > 0),
                  (this->get_n()/bn) + (this->get_n()%bn > 0),
                  0, 0 );

      void *buffer = A.get_buffer();
      Int_ cs = A.get_cs();

      Hier_ T;

      for (Int_ k2=0;k2<A.get_n();++k2) {
        for (Int_ k1=0;k1<A.get_m();++k1) {
          Int_ offm = k1*bm, offn = k2*bn;
          Int_ m = ( bm < (this->get_m()-offm) ? bm : this->get_m()-offm );
          Int_ n = ( bn < (this->get_n()-offn) ? bn : this->get_n()-offn );
          this->extract(T, m, n, offm, offn);
          elem<FLA_Obj>(buffer, cs, 1, k1, k2) = ~(T.flat());
        }
      }
    }

    Flat_ flat() {
      FLA_Obj obj;
      Int_ m=0, n=0, offm=0, offn=0;

      void *buffer = this->get_buffer();
      Int_ cs = this->get_cs();

      if (this->get_n()>0) 
        for (Int_ i=0;i<this->get_m();++i) 
          m+= elem<FLA_Obj>(buffer, cs, 1, i, 0).m;

      if (this->get_m()>0)
        for (Int_ j=0;j<this->get_n();++j)
          n+= elem<FLA_Obj>(buffer, cs, 1, 0, j).n;

      if (m && n) { 
        obj = elem<FLA_Obj>(buffer, cs, 1, 0, 0);
        obj.m    = m; obj.n    = n; 
      } else {
        obj.offm = 0; obj.offn = 0;
        obj.m    = m; obj.n    = n; 
      }
      return Flat_(obj);
    }

    virtual std::ostream& content(std::ostream& os) const {
      this->Matrix_::content(os);
      if (!this->is_created()) 
        return (os << "Buffer is null " << std::endl);

      Int_ cs = this->get_cs(), m = this->get_m(), n = this->get_n();
      void *buffer = this->get_buffer();

      for (Int_ k1=0;k1<this->get_m();++k1) {
        for (Int_ k2=0;k2<this->get_n();++k2) 
          fla_view(os, elem<FLA_Obj>(buffer, cs, 1, k1, k2));
        os << std::endl;
      }
      return (os << std::endl);
    }
  };
}

#endif
